import secrets
import hashlib
from publickey.math_utils import mod_inverse

# ==============================================================================
# 1. 曲线参数定义
# ==============================================================================

# --- secp256k1 参数 (比特币和许多其他加密货币使用的曲线) ---
# y^2 = x^3 + 7 (mod p)
_p_secp256k1 = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F
_a_secp256k1 = 0
_b_secp256k1 = 7
_Gx_secp256k1 = 0x79BE667EF9DCBBAC55A06295CE870B07029BFCDB2DCE28D959F2815B16F81798
_Gy_secp256k1 = 0x483ADA7726A3C4655DA4FBFC0E1108A8FD17B448A68554199C47D08FFB10D4B8
_n_secp256k1 = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141
_h_secp256k1 = 1

# --- 新增: secp256r1 (NIST P-256, prime256v1) 参数 ---
# 这是在TLS/SSL和许多Web标准中广泛使用的曲线
_p_secp256r1 = 0xFFFFFFFF00000001000000000000000000000000FFFFFFFFFFFFFFFFFFFFFFFF
_a_secp256r1 = 0xFFFFFFFF00000001000000000000000000000000FFFFFFFFFFFFFFFFFFFFFFFC
_b_secp256r1 = 0x5AC635D8AA3A93E7B3EBBD55769886BC651D06B0CC53B0F63BCE3C3E27D2604B
_Gx_secp256r1 = 0x6B17D1F2E12C4247F8BCE6E563A440F277037D812DEB33A0F4A13945D898C296
_Gy_secp256r1 = 0x4FE342E2FE1A7F9B8EE7EB4A7C0F9E162BCE33576B315ECECBB6406837BF51F5
_n_secp256r1 = 0xFFFFFFFF00000000FFFFFFFFFFFFFFFFBCE6FAADA7179E84F3B9CAC2FC632551
_h_secp256r1 = 1

# --- 新增: secp384r1 (NIST P-384) 参数 ---
# 提供比256位曲线更高的安全级别
_p_secp384r1 = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFFFF0000000000000000FFFFFFFF
_a_secp384r1 = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFFFF0000000000000000FFFFFFFC
_b_secp384r1 = 0xB3312FA7E23EE7E4988E056BE3F82D19181D9C6EFE8141120314088F5013875AC656398D8A2ED19D2A85C8EDD3EC2AEF
_Gx_secp384r1 = 0xAA87CA22BE8B05378EB1C71EF320AD746E1D3B628BA79B9859F741E082542A385502F25DBF55296C3A545E3872760AB7
_Gy_secp384r1 = 0x3617DE4A96262C6F5D9E98BF9292DC29F8F41DBD289A147CE9DA3113B5F0B8C00A60B1CE1D7E819D7A431D7C90EA0E5F
_n_secp384r1 = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFC7634D81F4372DDF581A0DB248B0A77AECEC196ACCC52973
_h_secp384r1 = 1


# ==============================================================================
# 2. 曲线注册表和获取函数
# ==============================================================================

predefined_curves = {
    "secp256k1": (_a_secp256k1, _b_secp256k1, _p_secp256k1, _Gx_secp256k1, _Gy_secp256k1, _n_secp256k1, _h_secp256k1),
    "secp256r1": (_a_secp256r1, _b_secp256r1, _p_secp256r1, _Gx_secp256r1, _Gy_secp256r1, _n_secp256r1, _h_secp256r1),
    "secp384r1": (_a_secp384r1, _b_secp384r1, _p_secp384r1, _Gx_secp384r1, _Gy_secp384r1, _n_secp384r1, _h_secp384r1),
}

def get_curve_by_name(name: str):
    """
    通过标准名称获取一个预定义的椭圆曲线实例。
    :param name: 曲线名称, e.g., "secp256k1", "secp256r1", "secp384r1"
    :return: 一个 EllipticCurve 对象
    """
    params = predefined_curves.get(name)
    if params is None:
        raise ValueError(f"未知的曲线名称: {name}。可用曲线: {list(predefined_curves.keys())}")
    return EllipticCurve(*params)


# ==============================================================================
# 3. 核心类和函数定义 (与之前基本相同)
# ==============================================================================

class EllipticCurve:
    # ... 此类定义内部无随机数生成，保持不变 ...
    def __init__(self, a, b, p, Gx, Gy, n, h=1):
        self.a = a
        self.b = b
        self.p = p
        self.G = (Gx, Gy)
        self.n = n
        self.h = h
        self.O = None 
        if not self.is_on_curve(self.G):
            raise ValueError("基点 G 不在曲线上")

    def is_on_curve(self, P):
        if P == self.O: return True
        x, y = P
        return (y * y - (x * x * x + self.a * x + self.b)) % self.p == 0

    def add_points(self, P, Q):
        if P == self.O: return Q
        if Q == self.O: return P
        x1, y1 = P; x2, y2 = Q
        if x1 == x2 and y1 != y2: return self.O
        if P == Q:
            if y1 == 0: return self.O
            m = ((3 * x1 * x1 + self.a) * mod_inverse(2 * y1, self.p)) % self.p
        else:
            if x1 == x2: return self.O
            m = ((y2 - y1) * mod_inverse(x2 - x1, self.p)) % self.p
        x3 = (m * m - x1 - x2) % self.p
        y3 = (m * (x1 - x3) - y1) % self.p
        return (x3, y3)

    def scalar_multiply(self, k, P):
        if k == 0 or P == self.O: return self.O
        k = k % self.n
        if k == 0: return self.O
        result, current_add = self.O, P
        if k < 0:
            k = -k
            current_add = (P[0], -P[1] % self.p)
        while k > 0:
            if k & 1: result = self.add_points(result, current_add)
            current_add = self.add_points(current_add, current_add)
            k >>= 1
        return result

def generate_keypair(curve: EllipticCurve):
    private_key = secrets.randbelow(curve.n - 1) + 1
    public_key_point = curve.scalar_multiply(private_key, curve.G)
    return private_key, public_key_point

def ecc_encrypt_simplified(public_key_point, curve: EllipticCurve, message: bytes):
    k_e = secrets.randbelow(curve.n - 1) + 1
    R = curve.scalar_multiply(k_e, curve.G)
    S_point = curve.scalar_multiply(k_e, public_key_point)
    if S_point == curve.O:
        raise Exception("计算出的共享密钥点是无穷远点，加密失败，请重试。")
    shared_secret_material = S_point[0].to_bytes((S_point[0].bit_length() + 7) // 8, 'big')
    symmetric_key = hashlib.sha256(shared_secret_material).digest()
    return R, bytes(c ^ symmetric_key[i % len(symmetric_key)] for i, c in enumerate(message))

def ecc_decrypt_simplified(private_key, curve: EllipticCurve, R_point, ciphertext_bytes: bytes):
    S_point = curve.scalar_multiply(private_key, R_point)
    if S_point == curve.O:
        raise Exception("计算出的共享密钥点是无穷远点，解密失败。")
    shared_secret_material = S_point[0].to_bytes((S_point[0].bit_length() + 7) // 8, 'big')
    symmetric_key = hashlib.sha256(shared_secret_material).digest()
    return bytes(c ^ symmetric_key[i % len(symmetric_key)] for i, c in enumerate(ciphertext_bytes))


# ==============================================================================
# 4. 主函数演示与测试
# ==============================================================================

if __name__ == '__main__':
    
    # 演示如何使用不同的曲线
    curve_names_to_test = ["secp256k1", "secp256r1", "secp384r1"]
    
    for curve_name in curve_names_to_test:
        print("\n" + "#"*20, f" 测试曲线: {curve_name} ", "#"*20)
        
        # 使用新函数按名称获取曲线实例
        try:
            curve = get_curve_by_name(curve_name)
        except ValueError as e:
            print(e)
            continue

        # 1. 密钥生成
        print("\n--- 密钥生成 ---")
        private_key, public_key = generate_keypair(curve)
        print(f"私钥 (d): {hex(private_key)}")
        print(f"公钥 (Q): X={hex(public_key[0])}, Y={hex(public_key[1])}")

        # 2. 消息加密
        print("\n--- 加密 ---")
        message = f"Secret message for {curve_name}!".encode('utf-8')
        print(f"原始消息: {message.decode()}")

        R_point, ciphertext = ecc_encrypt_simplified(public_key, curve, message)
        print(f"加密 R 点: X={hex(R_point[0])[:20]}..., Y={hex(R_point[1])[:20]}...")
        print(f"加密密文: {ciphertext.hex()}")

        # 3. 消息解密
        print("\n--- 解密 ---")
        decrypted_message = ecc_decrypt_simplified(private_key, curve, R_point, ciphertext)
        print(f"解密后的消息: {decrypted_message.decode()}")

        assert decrypted_message == message, "解密失败！"
        print("\n加解密成功！")